// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <optional>
#include <pollux/core/expressions.h>
#include <pollux/core/plan_node.h>
#include <pollux/substrait/pollux_substrait_signature.h>
#include <pollux/substrait/proto/substrait/algebra.pb.h>
#include <pollux/substrait/proto/substrait/plan.pb.h>
#include <pollux/type/type.h>

namespace kumo::pollux::substrait {
    struct ExtensionFunctionId {
        /// Substrait extension YAML file uri.
        std::string uri;

        /// Substrait signature used in the function extension declaration is a
        /// combination of the name of the function along with a list of input
        /// argument types.The format is as follows : <function
        /// name>:<short_arg_type0>_<short_arg_type1>_..._<short_arg_typeN> for more
        /// detail information about the argument type please refer to link
        /// https://substrait.io/extensions/#function-signature-compound-names.
        std::string signature;

        bool operator==(const ExtensionFunctionId &other) const {
            return (uri == other.uri && signature == other.signature);
        }
    };

    /// Assigns unique IDs to function signatures using ExtensionFunctionId.
    class SubstraitExtensionCollector {
    public:
        SubstraitExtensionCollector();

        /// Given a scalar function name and argument types, return the functionId
        /// using ExtensionFunctionId.
        int getReferenceNumber(
            const std::string &functionName,
            const std::vector<TypePtr> &arguments);

        /// Given an aggregate function name and argument types and aggregation Step,
        /// return the functionId using ExtensionFunctionId.
        int getReferenceNumber(
            const std::string &functionName,
            const std::vector<TypePtr> &arguments,
            core::AggregationNode::Step aggregationStep);

        /// Add extension functions to Substrait plan.
        void addExtensionsToPlan(::substrait::Plan *plan) const;

    private:
        /// A bi-direction hash map to keep the relation between reference number and
        /// either function or type signature.
        template<class T>
        class BiDirectionHashMap {
        public:
            /// Add (key, value) pair if doesn't exist already, i.e. forwardMap doesn't
            /// contain the key and reverseMap doesn't contain the value.
            ///
            /// @return True if the values were added successfully. False, otherwise.
            bool putIfAbsent(const int &key, const T &value);

            const melon::F14FastMap<int, ExtensionFunctionId> forwardMap() const {
                return forwardMap_;
            }

            const melon::F14FastMap<T, int> &reverseMap() const {
                return reverseMap_;
            }

        private:
            melon::F14FastMap<int, T> forwardMap_;
            melon::F14FastMap<T, int> reverseMap_;
        };

        /// Assigns unique IDs to function signatures using ExtensionFunctionId.
        int getReferenceNumber(const ExtensionFunctionId &extensionFunctionId);

        int functionReferenceNumber = -1;
        std::shared_ptr<BiDirectionHashMap<ExtensionFunctionId> > extensionFunctions_;
    };

    using SubstraitExtensionCollectorPtr =
    std::shared_ptr<SubstraitExtensionCollector>;
} // namespace kumo::pollux::substrait

namespace std {
    /// Hash function of kumo::pollux::substrait::ExtensionFunctionId.
    template<>
    struct hash<kumo::pollux::substrait::ExtensionFunctionId> {
        size_t operator()(
            const kumo::pollux::substrait::ExtensionFunctionId &k) const {
            size_t val = hash<std::string>()(k.uri);
            val = val * 31 + hash<std::string>()(k.signature);
            return val;
        }
    };
} // namespace std
